#include <iostream>
#include <eigen3/Eigen/Eigen>
#include <EKF.h>
#include <cmath>

EKF::EKF()
{
    m_head = NULL;
    count = 0;
}
EKF::EKF(Eigen::Vector4d x0, Eigen::Matrix4d P0)
{
    m_head = new state;
    m_head->x_k = x0;
    m_head->P_k = P0;
    count = 0;
}

EKF::~EKF()
{
    Removeall();
}

void EKF::Removeall()
{// 清除堆空间数据
    state *p = m_head, *ptemp;
    ptemp = NULL;
    while (p)
    {
        ptemp = p;
        p = p->m_next;
        delete ptemp;
    }
}

void EKF::Getx_k(Eigen::Matrix4Xd &x)
{// 读出预测数值
    state *p = m_head;
    for (int i = 0; i < count; i++)
    {
        x.col(i) = p->x_k;
        p = p->m_next;
    }
}
void EKF::Pushz_k(Eigen::Matrix2Xd &z)
{// 写入量测数值
    state *p = m_head;
    int count_new = count + z.cols();
    for (int i = count; i < count_new; i++)
    {
        p->z_k = z.col(i);
        p->m_next = new state;
        p = p->m_next;
    }
    count = count_new;
}

void EKF::EKFAlgo()
{
    double T = 0.1, k_x = 0.01, k_y = 0.05, g = 9.8;
    Eigen::Vector4d x_;
    Eigen::Matrix4d P_;
    Eigen::Vector2d z_;
    Eigen::Matrix4d f_x;// f对x的偏导
    Eigen::Matrix2Xd h_x = Eigen::Matrix2Xd::Zero(2, 4);
    Eigen::Matrix4Xd f_w = Eigen::Matrix4Xd::Zero(4, 2);;
    Eigen::Matrix2d h_v;
    Eigen::Matrix2d w, v;// 噪声方差阵
    w <<    0.3, 0,
            0, 0.3;
    v <<    8, 0, 
            0, 0.001;
    f_x <<  1, T, 0, 0,
            0, 0, 0, 0,
            0, 0, 1, T,
            0, 0, 0, 0;
    f_w <<  0, 0,
            1, 0,   
            0, 0,
            0, 1;
    h_v <<  1, 0,
            0, 1;

    state *p = m_head, *ptemp;
    while (p->m_next)
    {
        ptemp = p->m_next;
        // 状态预测
        x_(0) = p->x_k(0) + T * p->x_k(1);
        x_(1) = p->x_k(1) - k_x * T * pow(p->x_k(1), 2);
        x_(2) = p->x_k(2) + T * p->x_k(3);
        x_(3) = p->x_k(3) + k_y * T * pow(p->x_k(3), 2) - g * T;
        z_(0) = sqrt(pow(x_(0), 2) + pow(x_(2), 2));
        z_(1) = atan2(x_(0), x_(2));
        // Jacobi矩阵获取
        f_x(1, 1) = 1 - 2 * k_x * p->x_k(1) *T;
        f_x(3, 3) = 1 + 2 * k_y * p->x_k(3) *T;
        h_x(0, 0) = x_(0)/z_(0);
        h_x(0, 2) = x_(2)/z_(0);
        h_x(1, 0) = x_(2)/pow(z_(0), 2);
        h_x(1, 2) = -x_(0)/pow(z_(0), 2);
        // 方差预测
        P_ = f_x * p->P_k * f_x.transpose() + f_w * w *f_w.transpose();
        // 增益更新
        ptemp->K_k = P_ * h_x.transpose() * (h_x * P_ * h_x.transpose() + h_v * v * h_v.transpose()).inverse();
        // 状态更新
        ptemp->x_k = x_ + ptemp->K_k * (ptemp->z_k - z_);
        ptemp->P_k = P_ - ptemp->K_k * h_x * P_;

        p = ptemp;
    }
    
}